/*
 * Description: 进程帮助类
 */
#include "proc_util.h"
#include "string_util.h"

namespace {
    typedef NTSTATUS(NTAPI* _NtQueryInformationProcess)(HANDLE ProcessHandle, DWORD ProcessInformationClass,
        PVOID ProcessInformation, ULONG ProcessInformationLength, PULONG ReturnLength);
    const UINT MAX_CMD_LEN = 1024;
}

void ProcUtil::GetProcessIdByProcessName(const std::string& processName, std::vector<DWORD>& tokenVec)
{
    std::string name = StringUtil::UpperCase(processName);
    std::vector<PROCESSENTRY32> lst;
    if (GetProcessList(lst)) {
        for (auto it = lst.begin(); it != lst.end(); it++) {
            std::string upperName = StringUtil::UpperCase(it->szExeFile);
            if (upperName == name) {
                tokenVec.push_back(it->th32ProcessID);
            }
        }
    }
}

bool ProcUtil::GetProcessList(std::vector<PROCESSENTRY32>& lst)
{
    PROCESSENTRY32 pe32 = { 0 };
    HANDLE processSnapHandle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (processSnapHandle != INVALID_HANDLE_VALUE) {
        pe32.dwSize = sizeof(PROCESSENTRY32);
        BOOL ret = Process32First(processSnapHandle, &pe32);
        while (ret == TRUE) {
            lst.push_back(pe32);
            pe32.dwSize = sizeof(PROCESSENTRY32);
            ret = Process32Next(processSnapHandle, &pe32);
        }
        CloseHandle(processSnapHandle);
    }
    return !lst.empty();
}

bool ProcUtil::IsProcessExists(const std::string& processName)
{
    bool ret = false;
    std::string name = StringUtil::UpperCase(processName);
    std::vector<PROCESSENTRY32> lst;
    if (GetProcessList(lst) == TRUE) {
        for (auto it = lst.begin(); it != lst.end(); it++) {
            std::string tmpName = StringUtil::UpperCase(it->szExeFile);
            if (tmpName == name) {
                ret = true;
                break;
            }
        }
    }
    return ret;
}

std::string ProcUtil::GetProcessCommandLine(HANDLE processHandle)
{
    HMODULE module = nullptr;
    std::wstring commandLine(L"");
    do {
        module = LoadLibrary("Ntdll.dll");
        if (module == nullptr) {
            break;
        }
        _NtQueryInformationProcess pfnQuery =
            (_NtQueryInformationProcess)GetProcAddress(module, "NtQueryInformationProcess");
        if (pfnQuery == nullptr) {
            LOG_ERR("GetProcessCommandLine GetProcAddress Error: {}", GetLastError());
            break;
        }
        PROCESS_BASIC_INFORMATION pi = { 0 };
        NTSTATUS reNtStatus = pfnQuery(processHandle,
            ProcessBasicInformation,
            &pi,
            static_cast<ULONG>(sizeof(PROCESS_BASIC_INFORMATION)),
            nullptr);
        if (!NT_SUCCESS(reNtStatus)) {
            LOG_ERR("GetProcessCommandLine NtQueryInformationProcess Error: ", GetLastError());
            break;
        }
        PEB peb = { 0 };
        RTL_USER_PROCESS_PARAMETERS para = { 0 };
        if (ReadProcessMemory(processHandle, pi.PebBaseAddress, &peb, sizeof(peb), nullptr) == FALSE) {
            LOG_ERR("GetProcessCommandLine ReadProcessMemory Error1: ", GetLastError());
            break;
        }
        if (ReadProcessMemory(processHandle, peb.ProcessParameters, &para, sizeof(para), nullptr) == FALSE) {
            LOG_ERR("GetProcessCommandLine ReadProcessMemory Error2: ", GetLastError());
            break;
        }
        commandLine = GetProcessCommandLineParam(processHandle, para);
    } while (false);

    if (module != nullptr) {
        FreeLibrary(module);
    }
    return StringUtil::WStringToString(commandLine);
}

std::wstring ProcUtil::GetProcessCommandLineParam(HANDLE processHandle, const RTL_USER_PROCESS_PARAMETERS& para)
{
    WCHAR* getValue = new (std::nothrow) WCHAR[MAX_CMD_LEN];
    if (getValue == nullptr) {
        return L"";
    }
    SecureZeroMemory(getValue, MAX_CMD_LEN * sizeof(WCHAR));
    size_t length = 0;
    if (para.CommandLine.Length > (MAX_CMD_LEN - 1) * sizeof(WCHAR)) {
        length = (MAX_CMD_LEN - 1) * sizeof(WCHAR);
    } else {
        length = para.CommandLine.Length;
    }
    std::wstring commandLineParam;
    if (::ReadProcessMemory(processHandle, para.CommandLine.Buffer, getValue, length, nullptr) == TRUE) {
        std::wstring commandLine = getValue;
        commandLineParam = GetCommandLineParam(commandLine);
    }
    delete[] getValue;
    return commandLineParam;
}

std::wstring ProcUtil::GetCommandLineParam(const std::wstring& commandLine)
{
    std::wstring commonLineParam;
    if (commandLine.find('\"') != std::wstring::npos) {
        const WCHAR* p = commandLine.c_str();
        p += 1;
        WCHAR tmpPath[MAX_CMD_LEN] = { 0 };
        for (int i = 0; *p && *p != L'\"' && i < MAX_CMD_LEN - 1; p++) {
            tmpPath[i++] = *p;
        }
        commonLineParam = tmpPath;
    } else {
        commonLineParam = commandLine;
    }
    return commonLineParam;
}
